#include <time.h>
#include <iostream>
#include <cmath>
#include <complex>
#include <cstring>
#include <iomanip>
#include <mpi.h>
#define complex std::complex<double>

#include "../include/Gaussian.h"
#include "../include/Fourier.h"
#include "../include/NonlinearResponse.h"
#include "../include/MatrixAB.h"
#include "../include/CrankNicholson.h"
#include "../include/Constants.h"
#include "../include/Fileout.h"
#include "../include/Grid.h"
#include "../include/UnitConvert.h"
#include "../include/Absord.h"
#include "../include/Util.h"
#include "../include/Parameter.h"

int main(int argc, char **argv)
{
    //! ----------------Variable Assignment------------
    // Time--------------------------------------------
    const auto t0     = Parameter::t0;;
    const auto t_max  = Parameter::t_max;
    const auto dt     = Parameter::dt;
    const int  nt     = Parameter::nt;
    const int  fft_num= Parameter::fft_num;

    // Radial------------------------------------------
    const auto r0     = Parameter::r0;
    const auto dr_min = Parameter::dr_min;
    const auto dr_max = Parameter::dr_max;
    const int  nr     = Parameter::nr;
    const int  r_step_type = input::step_type;
    
    // Longitudinal------------------------------------
    const auto z0 = Parameter::z0;
    const auto dz = Parameter::dz;
    const int  nz = Parameter::nz;

    // nonlinear parameter-------------------------------
    const auto Kerr = Parameter::Kerr;
    const auto Ng = Parameter::Ng;
    
    int index_omega1 = Parameter::index_omega1;
    int index_omega2 = Parameter::index_omega2;
    //!-------------------Absorption Function---------------------------------
    // The thickness of the absorbing layer----------------
    int ab_len = 100;
    double ab_array[ab_len];
    Absord ab_function;
    ab_function(ab_array, ab_len);
    // frequency_array-------------------------------------
    double frequency_array[nt];
    Fourier::fftfreq(nt, frequency_array, dt);     // frequency_array
    double d_omega = 2. * PI * fabs(frequency_array[1] - frequency_array[0]);
    int index_THz = round((input::THz_max * Constants::omega1THz) / (d_omega));
    // fft_array--------------------------------------------
    complex fft_outc_E[nt]; // E_fft
    complex fft_outc_J[nt]; // J_fft
    complex fft_outc_P[nt]; // P_fft
    // NonliearResponse-------------------------------------
    double ionization_rate[nt];
    double electron_density[nt];
    double nonlinear_J[nt];
    double nonlinear_P[nt];

    //!---------------引入类-------------------------------------
    NonlinearResponse nonlinear(nt, dt, Kerr, Ng);
    CrankNicholson solve(nr);
    //!---------------MPI Initialization---------------------
    int rank, processors;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &processors);
    // nr_mpi
    int rows_per_proc_nr = nr / processors;
    int rows_extral_nr = nr % processors;
    int rows_local_nr = rows_per_proc_nr + (rank < rows_extral_nr ? 1 : 0);
    // nt_mpi
    int rows_per_proc_nt = fft_num / processors;
    int rows_extral_nt = fft_num % processors;
    int rows_local_nt = rows_per_proc_nt + (rank < rows_extral_nt ? 1 : 0);
    //!-----------------主线程 rank = 0--------------------------
    if (rank == 0)
    {
        double start = MPI_Wtime();
        cout <<"local rows_nr of rank 0 = "<< rows_local_nr << endl; 
        cout <<"local rows_nt of rank 0 = "<< rows_local_nt << endl; 
        // input laser------------------------------------
        auto Phase1 = Parameter::phase1; // cep
        auto Phase2 = Parameter::phase2; // cep
        auto Phase0 = Parameter::phase0; // Relative Phase
        auto FWHM = Parameter::fwhm;

        // input laser1------------------------------------
        auto Wavelength1 = Parameter::Wavelength1;
        auto Waist1 = Parameter::Waist1;
        auto Focus1 = Parameter::Focus1;
        auto E_peak1 = Parameter::E_peak1;

        // input laser2--------------------------------------
        auto Wavelength2 = Parameter::Wavelength2;
        auto Waist2 = Parameter::Waist2;
        auto Focus2 = Parameter::Focus2;
        auto E_peak2 = Parameter::E_peak2;

        //---------------------------------
        int nr_out = Parameter::nr_out;

        cout << "nr_out = " << nr_out << endl;
        cout << "nr = " << nr << endl;
        cout << "nz = " << nz << endl;

        //!-----------------"Class Initialization"--------------------
        Step step(r_step_type, dr_min, dr_max, nr);
        GaussianDual pulse(z0, FWHM, t_max, Phase0,
                           Waist1, Focus1, E_peak1, Wavelength1, Phase1,
                           Waist2, Focus2, E_peak2, Wavelength2, Phase2);
        MatrixAB matrixab(nr, fft_num, dz);
       
        //! -----------------------1D_Array-------------------------
        // Grid Generation----------------------------------------
        double *r_step_array = new double[nr]; // StepArray
        double *R_grid = new double[nr];       // RadialGrid
        double *T_grid = new double[nt];       // TimeGrid
        double *Z_grid = new double[nz];       // LongitudinalGid

        // Coefficient Matrix---------------------------------------
        double *A0 = new double[nr];
        double *B0 = new double[nr];
        double *A2 = new double[nr];
        double *B2 = new double[nr];

        complex **A1 = new complex *[fft_num];
        complex **B1 = new complex *[fft_num];
        for (int i = 0; i < fft_num; i++)
        {
            A1[i] = new complex[nr];
            B1[i] = new complex[nr];
        }
        
        // 1D_Array, Assignment----------------------------------------
        step(r_step_array);
        Grid::Grid_step(R_grid, nr, r0, r_step_array); // R_grid
        Grid::Grid_1d(T_grid, nt, dt, t0);             // T_grid
        Grid::Grid_1d(Z_grid, nz, dz, z0);             // Z_grid
        
        // A1,B1, Assignment----------------------------------------
        matrixab(frequency_array,
                 R_grid, r_step_array,
                 A0, B0, A1, B1, A2, B2); // Coefficient Matrix

        //! --------------------2D_Array----------------------------------------
        Grid::Grid_2d *grid_2d = new Grid::Grid_2d(nr, nt);
        // E_n0 initialization--------------------------------------
        Grid::ini_zeros(nr, nt, grid_2d->E_n0);
        Grid::ini_zeros(fft_num, nr, grid_2d->E_n0fft);
        Grid::ini_zeros(fft_num, nr, grid_2d->G_n0fft);
        Grid::ini_zeros(fft_num, nr, grid_2d->G_n_1fft);
        Grid::ini_zeros(nr, fft_num, grid_2d->E_n1fft);
        pulse(nr, nt, R_grid, T_grid, grid_2d->E_n0);

        //!---------------------Output Folder--------------------------------------
        Fileout::fileout_grid(Parameter::filename_grid, frequency_array, T_grid,
                              R_grid, Z_grid, ab_array, nt, nr, nz, ab_len);
        Parameter::fileout_parameter();
        Fileout::output_Erz_accuracy(Parameter::filename_Erz_accuracy,
                                     pulse, R_grid, Z_grid, nr_out, nz);

        //!----------------Release Space before loop-------------------------------------
        delete[] Z_grid;
        delete[] T_grid;
        delete[] R_grid;
        delete[] r_step_array;
        Z_grid = nullptr;
        T_grid = nullptr;
        R_grid = nullptr;
        r_step_array = nullptr;

        //!----------------------nonlinear array---------------------------
        double intensity_array[nr];
        double intensityTHz_array[nr];
        double plasmaDensity_array[nr];
        double phase1_array[nr];
        double phase2_array[nr];
        //!----------------------MPI array--------------------------------
        //每一个线程要处理的数据的行数
        int rows_proc_nr_array[processors];
        int rows_proc_nt_array[processors];
        if(processors > 1){
            for(int index = 1; index < processors; index++){
                int rows_proc_nr = rows_per_proc_nr + (index < rows_extral_nr ? 1 : 0);
                int rows_proc_nt = rows_per_proc_nt + (index < rows_extral_nt ? 1 : 0);
                rows_proc_nr_array[index] = rows_proc_nr;
                rows_proc_nt_array[index] = rows_proc_nt;
            }
        }
        //! ---------------send coefficient matrix to all processors---------------
        if (processors > 1){
            double start1 = MPI_Wtime();

            int rows_start = rows_local_nt;
            for(int index = 1; index < processors; index++){
                int rows_proc_nt = rows_proc_nt_array[index];
                int data_num0 = rows_proc_nt * nr;
                complex* send_A1 = new complex[data_num0];
                complex* send_B1 = new complex[data_num0];

                Util::vectorization(A1, send_A1, rows_proc_nt, nr, rows_start);
                Util::vectorization(B1, send_B1, rows_proc_nt, nr, rows_start);

                MPI_Send(A0, nr, MPI_DOUBLE, index, 0, MPI_COMM_WORLD);
                MPI_Send(B0, nr, MPI_DOUBLE, index, 1, MPI_COMM_WORLD);
                MPI_Send(A2, nr, MPI_DOUBLE, index, 2, MPI_COMM_WORLD);
                MPI_Send(B2, nr, MPI_DOUBLE, index, 3, MPI_COMM_WORLD);
                MPI_Send(send_A1, data_num0, MPI_C_DOUBLE_COMPLEX, index, 4, MPI_COMM_WORLD);
                MPI_Send(send_B1, data_num0, MPI_C_DOUBLE_COMPLEX, index, 5, MPI_COMM_WORLD);

                rows_start += rows_proc_nt;
                delete[] send_A1;
                delete[] send_B1;
                send_A1 = nullptr;
                send_B1 = nullptr;
            }

            double end1 = MPI_Wtime();
            double time_s1 = end1 - start1;
            cout << "send coefficient matrix to all processors time: " << time_s1 << endl;
        }
//!-------------------into outer loop nz--------------------------------------------
        cout << "--------into outer loop nz--------" << endl;       
        for (int num_z = 0; num_z < nz; num_z++)
        {

            std::cout << "nz=" << num_z << std::endl;
            //!----------------fft-openmpi-start---------------------------
            //主线程分发数据 fft
            if(processors > 1){
                int rows_start_nr = rows_local_nr;
                for (int index = 1; index < processors; index++){
                    int rows_proc_nr = rows_proc_nr_array[index];
                    int data_num1 = rows_proc_nr * nt;
                    double* send_En0 = new double[data_num1];
                    Util::vectorization(grid_2d->E_n0, send_En0, rows_proc_nr, nt, rows_start_nr);
                    MPI_Send(send_En0, data_num1, MPI_DOUBLE, index, index + 6, MPI_COMM_WORLD);
                    rows_start_nr += rows_proc_nr;

                    delete[] send_En0;
                    send_En0 = nullptr; 
                }
            }
        
            //主线程执行任务 fft
            for(int num_r = 0; num_r < rows_local_nr; num_r++){

                int num_ab = (num_r < nr - ab_len)? 0 : num_r - nr + ab_len;
                // nonlinear response
                nonlinear(ionization_rate, electron_density,
                          nonlinear_J, nonlinear_P,
                          grid_2d->E_n0[num_r]);
                // fft
                Fourier::fft_r2c(nt, grid_2d->E_n0[num_r], fft_outc_E);
                Fourier::fft_r2c(nt, nonlinear_J, fft_outc_J);
                Fourier::fft_r2c(nt, nonlinear_P, fft_outc_P);

                complex iwJ = 0;
                complex w2P = 0;

                double intensity_all = 0;
                double intensity_THz = 0;
                double plasmaDensity = 0;

                for (int num_t = 0; num_t < fft_num; num_t++)
                {
                    double ab_i = ab_array[num_ab];
                    double omega_i = 2. * PI * frequency_array[num_t];
                    complex En0_fft = fft_outc_E[num_t];
                    grid_2d->E_n0fft[num_t][num_r] = En0_fft * (num_r < nr - ab_len ? 1.0 : ab_i); // 吸收边

                    iwJ = fft_outc_J[num_t] * omega_i * Constants::Ii;
                    w2P = -fft_outc_P[num_t] * omega_i * omega_i;
                    complex Gn0_fft = dz * (iwJ + w2P);  
                    grid_2d->G_n0fft[num_t][num_r] = Gn0_fft * (num_r < nr - ab_len ? 1.0 : ab_i); // 吸收边
            
                    intensity_all += fabs(grid_2d->E_n0fft[num_t][num_r]);
                    intensity_THz = (num_t <= index_THz) ? intensity_all : intensity_THz;
                }

                intensity_all = intensity_all / ((nt * 1.) / 2.);
                intensity_THz = intensity_THz / ((nt * 1.) / 2.);
                plasmaDensity = electron_density[nt - 1];

                intensity_array[num_r] = intensity_all;
                intensityTHz_array[num_r] = intensity_THz;
                plasmaDensity_array[num_r] = plasmaDensity;
            
                if (num_r == 0)
                {
                    Fileout::output_Ewzc(Parameter::THzSpectrumyOut, nt, fft_outc_E);
                    std::cout << "intensity: " << intensity_all << std::endl;
                }   
            }

            //主线程接收数据 fft
            if(processors > 1){
                int rows_start_nr = rows_local_nr;
                for (int index = 1; index < processors; index++){
                    int rows_proc_nr = rows_proc_nr_array[index];
                    int data_num2 = rows_proc_nr * fft_num;
                    //定义接收数组
                    complex* recv_En0fft = new complex[data_num2];
                    complex* recv_Gn0fft = new complex[data_num2];

                    double recv_intensity_array[rows_proc_nr];
                    double recv_intensityTHz_array[rows_proc_nr];
                    double recv_plasmaDensity_array[rows_proc_nr];
                    
                    //接收
                    MPI_Recv(recv_En0fft, data_num2, MPI_C_DOUBLE_COMPLEX, index, index + 7, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    MPI_Recv(recv_Gn0fft, data_num2, MPI_C_DOUBLE_COMPLEX, index, index + 8, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    
                    MPI_Recv(recv_intensity_array, rows_proc_nr,    MPI_DOUBLE, index, index + 9, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    MPI_Recv(recv_intensityTHz_array, rows_proc_nr, MPI_DOUBLE, index, index + 10, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    MPI_Recv(recv_plasmaDensity_array, rows_proc_nr,MPI_DOUBLE, index, index + 11, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    //copy data
                    Util::trans_reshape(grid_2d->E_n0fft, recv_En0fft, fft_num, rows_proc_nr, rows_start_nr);
                    Util::trans_reshape(grid_2d->G_n0fft, recv_Gn0fft, fft_num, rows_proc_nr, rows_start_nr);
                    for(int j = 0; j < rows_proc_nr; j++){
                        int j_index = rows_start_nr + j;    
                        intensity_array    [j_index] = recv_intensity_array[j];
                        intensityTHz_array [j_index] = recv_intensityTHz_array[j];
                        plasmaDensity_array[j_index] = recv_plasmaDensity_array[j];
                    }
                    rows_start_nr += rows_proc_nr;

                    delete[] recv_En0fft;
                    delete[] recv_Gn0fft;
                    recv_En0fft = nullptr;
                    recv_Gn0fft = nullptr;
                }
            }
            //!-------------------fft-openmpi-end----------------------------------

            //!--------------------------处理接收后的数据---------------------------
            for(int num_r = 0; num_r < nr; num_r++){
                for(int num_t = 0; num_t < fft_num; num_t ++){
                    complex Gn0_fft = grid_2d->G_n0fft[num_t][num_r];
                    complex Gn_1 = (num_z == 0) ? Gn0_fft : grid_2d->G_n_1fft[num_t][num_r];
                    grid_2d->G_n0fft[num_t][num_r] = 3. * Gn0_fft - Gn_1;  // Adams-Bashforth差分格式
                    grid_2d->G_n_1fft[num_t][num_r] = Gn0_fft;
                }
                phase1_array[num_r] = Util::Get_phase(grid_2d->E_n0fft[index_omega1][num_r]);
                phase2_array[num_r] = Util::Get_phase(grid_2d->E_n0fft[index_omega2][num_r]);  
            }
            
            //!-------------------------------输出数据----------------------------------
            Fileout::output_Etz(Parameter::Etz, nt, grid_2d->E_n0[0]);
            Fileout::fileout_1d_array(Parameter::intensity, intensity_array, nr_out, 1.);
            Fileout::fileout_1d_array(Parameter::intensity_r0, intensity_array, 1, 1.);

            Fileout::fileout_1d_array(Parameter::THzIntensity, intensityTHz_array, nr_out, 1.);
            Fileout::fileout_1d_array(Parameter::THzIntensity_r0, intensityTHz_array, 1, 1.);

            Fileout::fileout_1d_array(Parameter::plasmaDensity, plasmaDensity_array, nr_out, 1.);
            Fileout::fileout_1d_array(Parameter::plasmaDensity_r0, plasmaDensity_array, 1, 1.);

            Fileout::fileout_1d_array(Parameter::out_phase1, phase1_array, nr_out, 1.);
            Fileout::fileout_1d_array(Parameter::out_phase1_r0, phase1_array, 1, 1.);
            if (E_peak2 != 0)
            {
                Fileout::fileout_1d_array(Parameter::out_phase2, phase2_array, nr_out, 1.);
                Fileout::fileout_1d_array(Parameter::out_phase2_r0, phase2_array, 1, 1.);

                Fileout::outputErzaDual(Parameter::Erza, nr_out, grid_2d->E_n0fft[index_omega1], grid_2d->E_n0fft[index_omega2]);
                Fileout::outputErzaDual(Parameter::Erza_r0, 0, grid_2d->E_n0fft[index_omega1], grid_2d->E_n0fft[index_omega2]);
                Fileout::output_Erza(Parameter::Erza1, nr_out, grid_2d->E_n0fft[index_omega1]);
                Fileout::output_Erza(Parameter::Erza2, nr_out, grid_2d->E_n0fft[index_omega2]);
                Fileout::output_Erza(Parameter::Erza1_r0, 1, grid_2d->E_n0fft[index_omega1]);
                Fileout::output_Erza(Parameter::Erza2_r0, 0, grid_2d->E_n0fft[index_omega2]);
            }
            else if (E_peak2 == 0)
            {
                Fileout::output_Erza(Parameter::Erza, nr_out, grid_2d->E_n0fft[index_omega1]);
                Fileout::output_Erza(Parameter::Erza_r0, 1, grid_2d->E_n0fft[index_omega1]);
            }
            
            //!--------------solve-openmpi-start------------------------
            // 主线程分发数据 solve
            if(processors > 1){
                int rows_start_nt = rows_local_nt;
                for(int index = 1; index < processors; index++){
                    int rows_proc_nt = rows_proc_nt_array[index];
                    int data_num0 = rows_proc_nt * nr;
                    complex* send_En0fft = new complex[data_num0];
                    complex* send_Gn0fft = new complex[data_num0];

                    Util::vectorization(grid_2d->E_n0fft, send_En0fft, rows_proc_nt, nr, rows_start_nt);
                    Util::vectorization(grid_2d->G_n0fft, send_Gn0fft, rows_proc_nt, nr, rows_start_nt);
                    
                    MPI_Send(send_En0fft, data_num0, MPI_C_DOUBLE_COMPLEX, index, index + 12, MPI_COMM_WORLD);
                    MPI_Send(send_Gn0fft, data_num0, MPI_C_DOUBLE_COMPLEX, index, index + 13, MPI_COMM_WORLD);
                    rows_start_nt += rows_proc_nt;

                    delete[] send_En0fft;
                    delete[] send_Gn0fft;
                    send_En0fft = nullptr;
                    send_Gn0fft = nullptr;
                }
            }

            // 主线程执行任务 solve
            complex X1[nr];
            for (int num_t = 0; num_t < rows_local_nt; num_t++)
            {
                // solve
                solve(A0, A1[num_t], A2,
                   B0, B1[num_t], B2,
                   grid_2d->E_n0fft[num_t], grid_2d->G_n0fft[num_t], X1);

                for (int num_r = 0; num_r < nr; num_r++)
                {
                    grid_2d->E_n1fft[num_r][num_t] = X1[num_r];
                    // 对num_t行数据求解，将得到的解放在E_n1fft 的第num_r行 num_t列
                }
            }
            // 主线程接收数据 solve
            if(processors > 1){
                int rows_start_nt = rows_local_nt;
                for (int index = 1; index < processors; index++){
                    int rows_proc_nt = rows_proc_nt_array[index];
                    int data_num0 = rows_proc_nt * nr;
                    complex* recv_En1fft = new complex[data_num0];

                    MPI_Recv(recv_En1fft, data_num0, MPI_C_DOUBLE_COMPLEX, index, index + 14, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    
                    Util::trans_reshape(grid_2d->E_n1fft, recv_En1fft, nr, rows_proc_nt, rows_start_nt);
                    rows_start_nt += rows_proc_nt;

                    delete[] recv_En1fft;
                    recv_En1fft = nullptr;
                }
            }
            //!-----------------solve-openmpi-end------------------------

            //!-----------------ifft—openmpi-start-------------------------
            // 主线程分发数据 ifft
            if(processors > 1){
                int rows_start_nr = rows_local_nr;
                for (int index = 1; index < processors; index++){
                    int rows_proc_nr = rows_proc_nr_array[index];
                    int data_num2 = rows_proc_nr * fft_num;
                    complex* send_En1fft = new complex[data_num2];
                    Util::vectorization(grid_2d->E_n1fft, send_En1fft, rows_proc_nr, fft_num, rows_start_nr);
                    
                    MPI_Send(send_En1fft, data_num2, MPI_C_DOUBLE_COMPLEX, index, index + 15, MPI_COMM_WORLD);    
                    rows_start_nr += rows_proc_nr;

                    delete[] send_En1fft;
                    send_En1fft = nullptr;
                }
            }
            // 主线程执行任务 ifft
            for (int num_r = 0; num_r < rows_local_nr; num_r++)
            {
                //ifft
                Fourier::ifft_c2r(nt, grid_2d->E_n1fft[num_r], grid_2d->E_n0[num_r]);
            }

            // 主线程接收数据 ifft
            if(processors > 1){
                int rows_start_nr = rows_local_nr;
                for (int index = 1; index < processors; index++){
                    int rows_proc_nr = rows_proc_nr_array[index];
                    int data_num2 = rows_proc_nr * nt;
                    double* recv_En1 = new double[data_num2];
                    MPI_Recv(recv_En1, data_num2, MPI_DOUBLE, index, index + 16, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                    Util::reshape(grid_2d->E_n0, recv_En1, rows_proc_nr, nt, rows_start_nr);
                    rows_start_nr += rows_proc_nr;
                    
                    delete[] recv_En1;
                    recv_En1 = nullptr;
                }
            }
            //!-----------------ifft—openmpi-end---------------------------
        }

        // Release Space----------------------------------------
        delete grid_2d;

        delete[] A0;
        delete[] A2;
        delete[] B0;
        delete[] B2;
        for (int i = 0; i < fft_num; i++)
        {
            delete[] A1[i];
            A1[i] = nullptr;
            delete[] B1[i];
            B1[i] = nullptr;
        }
        delete[] A1;
        A1 = nullptr;
        delete[] B1;
        B1 = nullptr;

        //print time-----------------------------------------------------------------------
        double end = MPI_Wtime();
        double time_s = end - start;
        int time_h = int(time_s) / 3600;
        int time_min = int(time_s) % 3600 / 60;
        time_s = time_s - time_h * 3600. - time_min * 60.;
        printf("Time with %d process: %02d:%02d:%.1f\n", processors, time_h, time_min, time_s);
    }
    else{
        //!--------------接收系数矩阵----------------------
        int data_num0 = rows_local_nt * nr;
        double *proc_A0 = new double[nr];
        double *proc_B0 = new double[nr];
        double *proc_A2 = new double[nr];
        double *proc_B2 = new double[nr];
        complex *proc_A1 = new complex[data_num0];
        complex *proc_B1 = new complex[data_num0];
        MPI_Recv(proc_A0, nr, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Recv(proc_B0, nr, MPI_DOUBLE, 0, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Recv(proc_A2, nr, MPI_DOUBLE, 0, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Recv(proc_B2, nr, MPI_DOUBLE, 0, 3, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Recv(proc_A1, data_num0, MPI_DOUBLE_COMPLEX, 0, 4, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Recv(proc_B1, data_num0, MPI_DOUBLE_COMPLEX, 0, 5, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

        //!----------------fft-ifft数组-----------------
        int data_num1 = rows_local_nr * fft_num;
        int data_num2 = rows_local_nr * nt;
        // En_2d
        double *proc_En = new double[data_num2];
        // Enfft_2d,Gnfft_2d
        complex *proc_Enfft = new complex[data_num1];
        complex *proc_Gnfft = new complex[data_num1];

        // 定义fft-ifft暂存数组
        double send_intensity_array[rows_local_nr];
        double send_intensityTHz_array[rows_local_nr];
        double send_plasmaDensity_array[rows_local_nr];
        double  En_1d[nt];
        complex Enfft_1d[fft_num];

        // 定义用于接收solve的数组
        complex* proc_En0fft_solve = new complex[data_num0];
        complex* proc_Gn0fft_solve = new complex[data_num0];
        complex* proc_En1fft_solve = new complex[data_num0];
        // 定义solve暂存数据
        complex En0fft_solve_1d[nr];
        complex Gn0fft_solve_1d[nr];
        complex proc_A1_1d[nr];
        complex proc_B1_1d[nr];

        //----------------------nz循环----------------------------
        for (int num_z = 0; num_z < nz; num_z++)
        {
            //! -----------------fft并行-------------------------
            // 线程接收数据 fft
            MPI_Recv(proc_En, data_num2, MPI_DOUBLE, 0, rank + 6, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            // 线程处理数据 fft
            for (int num_r = 0; num_r < rows_local_nr; num_r++)
            {
                int num_ab = (num_r < nr - ab_len) ? 0 : num_r - nr + ab_len;

                int start_index = num_r * nt;
                Util::splid(proc_En, En_1d, start_index, nt);
                
                // nonlinear response
                nonlinear(ionization_rate, electron_density,
                          nonlinear_J, nonlinear_P,
                          En_1d);
                // fft
                Fourier::fft_r2c(nt, En_1d,  fft_outc_E);
                Fourier::fft_r2c(nt, nonlinear_J, fft_outc_J);
                Fourier::fft_r2c(nt, nonlinear_P, fft_outc_P);

                complex iwJ = 0;
                complex w2P = 0;

                double intensity_all = 0;
                double intensity_THz = 0;
                double plasmaDensity = 0;

                for (int num_t = 0; num_t < fft_num; num_t++)
                {
                    int num_t_index = num_r * fft_num + num_t;

                    double ab_i = ab_array[num_ab];
                    double omega_i = 2. * PI * frequency_array[num_t];

                    complex En0_fft = fft_outc_E[num_t];
                    proc_Enfft[num_t_index] = En0_fft * (num_r < nr - ab_len ? 1.0 : ab_i) ;

                    iwJ = fft_outc_J[num_t] * omega_i * Constants::Ii;
                    w2P = -fft_outc_P[num_t] * omega_i * omega_i;
                    complex Gn0_fft = dz * (iwJ + w2P);
                    proc_Gnfft[num_t_index] = Gn0_fft * (num_r < nr - ab_len ? 1.0 : ab_i);

                    intensity_all += fabs(proc_Enfft[num_t_index]);
                    intensity_THz = (num_t <= index_THz) ? intensity_all : intensity_THz;
                }
                intensity_all = intensity_all / ((nt * 1.) / 2.);
                intensity_THz = intensity_THz / ((nt * 1.) / 2.);
                plasmaDensity = electron_density[nt - 1];

                send_intensity_array[num_r] = intensity_all;
                send_intensityTHz_array[num_r] = intensity_THz;
                send_plasmaDensity_array[num_r] = plasmaDensity;
            }
            // 线程发送数据 fft 
            MPI_Send(proc_Enfft, data_num1, MPI_C_DOUBLE_COMPLEX, 0, rank + 7, MPI_COMM_WORLD);
            MPI_Send(proc_Gnfft, data_num1, MPI_C_DOUBLE_COMPLEX, 0, rank + 8, MPI_COMM_WORLD);

            MPI_Send(send_intensity_array,    rows_local_nr, MPI_DOUBLE, 0, rank + 9, MPI_COMM_WORLD);
            MPI_Send(send_intensityTHz_array, rows_local_nr, MPI_DOUBLE, 0, rank + 10, MPI_COMM_WORLD);
            MPI_Send(send_plasmaDensity_array,rows_local_nr, MPI_DOUBLE, 0, rank + 11, MPI_COMM_WORLD);


            //!-----------------solve并行-----------------------
            // 线程接收数据 solve
            MPI_Recv(proc_En0fft_solve, data_num0, MPI_C_DOUBLE_COMPLEX, 0, rank + 12, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Recv(proc_Gn0fft_solve, data_num0, MPI_C_DOUBLE_COMPLEX, 0, rank + 13, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            // 线程处理数据 solve
            complex X1[nr];
            for (int num_t = 0; num_t < rows_local_nt; num_t++)
            {
                int start_index = num_t * nr;
                Util::splid(proc_En0fft_solve, En0fft_solve_1d, start_index, nr);
                Util::splid(proc_Gn0fft_solve, Gn0fft_solve_1d, start_index, nr);
                Util::splid(proc_A1, proc_A1_1d, start_index, nr);
                Util::splid(proc_B1, proc_B1_1d, start_index, nr);

                // solve
                solve(proc_A0, proc_A1_1d, proc_A2,
                      proc_B0, proc_B1_1d, proc_B2,
                      En0fft_solve_1d, Gn0fft_solve_1d, X1);

                Util::combine(proc_En1fft_solve, X1, start_index, nr);
                
            }
            // 线程发送数据 solve
            MPI_Send(proc_En1fft_solve, data_num0, MPI_C_DOUBLE_COMPLEX, 0, rank + 14, MPI_COMM_WORLD);

            //!----------------ifft并行------------------------
            // 线程接收数据 ifft
            MPI_Recv(proc_Enfft, data_num1, MPI_C_DOUBLE_COMPLEX, 0, rank + 15, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            // 线程处理数据 ifft 
            for(int num_r = 0; num_r < rows_local_nr; num_r ++)
            {
                int start_index1 = num_r * fft_num;
                Util::splid(proc_Enfft, Enfft_1d, start_index1, fft_num);
                
                // ifft
                Fourier::ifft_c2r(nt, Enfft_1d, En_1d);

                int start_index2 = num_r * nt;
                Util::combine(proc_En, En_1d, start_index2, nt);
            }
            // 线程发送数据 ifft
            MPI_Send(proc_En, data_num2, MPI_DOUBLE, 0, rank + 16, MPI_COMM_WORLD);
        }
        delete[] proc_En;
        proc_En = nullptr;
        delete[] proc_Enfft;
        proc_Enfft = nullptr;
        delete[] proc_Gnfft;
        proc_Gnfft = nullptr;

        delete[] proc_En0fft_solve;
        proc_En0fft_solve = nullptr;
        delete[] proc_Gn0fft_solve;
        proc_Gn0fft_solve = nullptr;
        delete[] proc_En1fft_solve;
        proc_En1fft_solve = nullptr;

        delete[] proc_A0;
        proc_A0 = nullptr;
        delete[] proc_B0;
        proc_B0 = nullptr;
        delete[] proc_A1;
        proc_A1 = nullptr;
        delete[] proc_B1;
        proc_B1 = nullptr;
        delete[] proc_A2;
        proc_A2 = nullptr;
        delete[] proc_B2;
        proc_B2 = nullptr;
    }
    MPI_Finalize();
    return 0;
}